/*
    Copyright 2006-2011 Patrik Jonsson, sunrise@familjenjonsson.org

    This file is part of Sunrise.

    Sunrise is free software; you can redistribute it and/or modify
    it under the terms of the GNU General Public License as published by
    the Free Software Foundation; either version 3 of the License, or
    (at your option) any later version.

    Sunrise is distributed in the hope that it will be useful,
    but WITHOUT ANY WARRANTY; without even the implied warranty of
    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    GNU General Public License for more details.

    You should have received a copy of the GNU General Public License
    along with Sunrise.  If not, see <http://www.gnu.org/licenses/>.

*/

/// \file
/// FITS I/O functions for blitz arrays.  This file contains routines
/// to read/write columns from the FITS files into blitz arrays and
/// vice versa.  All of these are pretty dependent on the storage
/// formats of vectors, arrays and TinyVectors, so beware if you
/// change compiler/architecture...

// $Id$


#ifndef __blitz_fits__
#define __blitz_fits__

#include "blitz/array.h"


#include "CCfits/CCfits"
#include <iterator>
#include <vector>
#include <stdexcept>


inline void checkerr(int status)
{
  // and report any error and throw
  if (status) {
    fits_report_error (stderr, status);
    throw std::runtime_error("FITS error while reading");
  }
}


/** This "one-dimensional vector or array" class gives a unified
    interface for writing FITS columns of vectors or blitz arrays.  */
template<typename T>
class onedva {
public:
  virtual ~onedva() {};
  
  /// Writes the vector/array to a FITS column.
  virtual void write_column(CCfits::ExtHDU& hdu,
			    CCfits::Column& col,
			    int start_row, int start_element, int num) const=0;
};

/** Derived class implementing the onedva functionality for vector.  */
template<typename T>
class onedva_v : public onedva<T> {
private:
  const std::vector<T>& v;

public:
  onedva_v (const std::vector<T> & vv): v (vv) {};

  virtual void write_column(CCfits::ExtHDU& hdu,
			    CCfits::Column& col,
			    int start_row, int start_element, int num) const {
    assert (start_element +num <= v.size());
    int status=0;
    CCfits::FITSUtil::MatchType<T> imageType;
    // the const_cast is necessary because the cfitsio functions do
    // not distinguish between void pointers and const void
    // pointers. Fits_write_col does not modify the argument in
    // any way, so this is safe.
    fits_write_col (hdu.fitsPointer(), imageType(), col.index(), 
		    start_row,1,num, const_cast<T*>(&v[start_element]),
		    &status); 
    // should we throw an exception on failure here?
  }
};

/** Derived class implementing the onedva functionality for a vector
    range. This works just like a subvector was copied and used.  */
template<typename T_iterator>
class onedva_vr : 
  public onedva<typename std::iterator_traits<T_iterator>::value_type> {
private:
  T_iterator b,e;

public:
  onedva_vr (const T_iterator& bb, 
	     const T_iterator& ee): 
    b (bb), e (ee) {};

  virtual void write_column(CCfits::ExtHDU& hdu,
			    CCfits::Column& col,
			    int start_row, int start_element, int num) const {
    assert (start_element+num <= e-b);
    int status=0;
    CCfits::FITSUtil::MatchType<typename std::iterator_traits<T_iterator>::value_type> imageType;
    // the const_cast is necessary because the cfitsio functions do
    // not distinguish between void pointers and const void
    // pointers. Fits_write_col does not modify the argument in
    // any way, so this is safe.
    fits_write_col (hdu.fitsPointer(), imageType(), col.index(), 
		    start_row,1,num, 
		    const_cast<typename std::iterator_traits<T_iterator>::value_type*>(&*(b+start_element)),
		    &status); 
    // should we throw an exception on failure here?
  }
};

/** Derived class implementing the onedva functionality for array.  */
template<typename T>
class onedva_a : public onedva<T> {
private:
  const blitz::Array<T, 1>& a;

public:
  onedva_a (const blitz::Array<T, 1>& aa): a (aa) {};

  /** The write_column member forwards to the write() function. */
  virtual void write_column(CCfits::ExtHDU& hdu,
			    CCfits::Column& col,
			    int start_row, int start_element, int num) const {
    write (hdu.fitsPointer(), col.index(),
	   a ( blitz::Range(start_element, start_element+num-1 )), start_row);}
};

/// Factory function to easily create a onedva from a vector.
template<typename T>
const onedva_v<T> make_onedva(const std::vector<T>& v)
{
  return onedva_v<T>(v);
}

/// Factory function to easily create a onedva from a vector range.
template<typename T_iterator>
const onedva_vr<T_iterator>
make_onedva(T_iterator b, T_iterator e)
{
  return onedva_vr<T_iterator>(b,e);
}

/// Factory function to easily create a onedva from an array.
template<typename T>
const onedva_a<T> make_onedva(const blitz::Array<T, 1>& a)
{
  return onedva_a<T>(a);
}



/** Reads a CCfits table column into a one-dimensional blitz array. */
template < class T >
void read (CCfits::Column & c, blitz::Array<T,1> & a)
{
  using namespace blitz;
  std::vector <T> temporary;
  c.read(temporary, 1, c.rows() );
  // I'm not sure this is kosher but except for a very pathological
  //vector it should be safe. No you ARE allowed to do this with
  //std::vector
  const int n = temporary.size();
  try {
    a.resize(n);
  }
  catch (...) {
    std::cerr << "One-dimensional column read can't allocate output array!"  << std::endl;
    throw;
  }
  // must duplicate the data since otherwise we are using a dead
  // pointer
  a = Array<T, 1>(& temporary[0], shape(n), duplicateData);
}

/** Reads a CCfits table column into a two-dimensional blitz array. */
template < class T >
void read (CCfits::Column & c, blitz::Array<T,2> & a)
{
  using namespace blitz;
  c.parent()->makeThisCurrent();

  int status = 0;
  // Find dimensions
  long rows;
  fits_get_num_rows(c.parent()->fitsPointer(), &rows, &status);
  int type;
  long repeat, width;
  fits_get_coltype(c.parent()->fitsPointer(), c.index(),
		   &type, &repeat, &width, &status);
  // (type and width not really needed)
  
  // now we can resize the array
  a.resize(rows, repeat);
  CCfits::FITSUtil::MatchType<T> imageType;  

  if(a.isStorageContiguous()) {
    // we can read the data stright into the blitz array
    fits_read_col(c.parent()->fitsPointer(), imageType(), c.index(), 
		  1, 1, a.size(), 0, a.dataFirst(), 0, &status);
  }
  else {
    // a is not contiguous, we must read into a temporary
    Array<T,2> temp(rows, repeat, contiguousArray);
    fits_read_col(c.parent()->fitsPointer(), imageType(), c.index(), 
		  1, 1, temp.size(), 0, temp.dataFirst(), 0, &status);
    a=temp;
  }

  checkerr(status);    
}

/** Reads a CCfits table column into a vector of blitz TinyVectors. */
template < class T, int N >
void read (CCfits::Column & c, std::vector<blitz::TinyVector<T, N> >& v)
{
  if (c.repeat() != N)
    throw CCfits::Column::RangeError ("Repeat count not consistent with array size");
  std::vector <valarray <T> > temporary;
  c.readArrays(temporary, 1, c.rows() );
  const int n1 = temporary.size();
  assert (temporary [0].size() == N);
  v.clear();
  v.resize(n1);
  // copying the data is a little more painful this way, but there is
  // no assignment operator between valaarry and tinyvector
  for (int i = 0; i < n1; ++i)
    for (int j = 0; j < N; ++j)
      v [i] [j] = temporary [i] [j];
}

/** Writes a one-dimensional blitz array to a CCfits table vector column. */
template < class T >
void write (CCfits::Column & c, const blitz::Array<T,1>& array, long first_row)
{
  // ensure data is contiguous in memory
  blitz::Array<T,1> a (blitz::contiguousArray); 
  if (array.isStorageContiguous())
    a.reference(array);
  else {
    a.resize(array.shape());
    a=array;
    assert(a.isStorageContiguous());
  }

  c.write(a.dataFirst(), a.rows(), first_row );
}

/** Writes a one-dimensional blitz array to a cfitsio table vector
    column in the current HDU. */
template < class T >
void write (fitsfile* const fptr, int colnum, const blitz::Array<T,1>& array, 
	    long first_row)
{
  // ensure data is contiguous in memory
  blitz::Array<T,1> a (blitz::contiguousArray); 
  if (array.isStorageContiguous())
    a.reference(array);
  else {
    a.resize(array.shape());
    a=array;
    assert(a.isStorageContiguous());
  }

  int status=0;
  fits_write_col (fptr, TDOUBLE, colnum,
		  first_row,1,a.rows(), a.dataFirst(), &status);
  fits_report_error(stderr,status);
}

/** Writes a two-dimensional blitz array to a CCfits table vector column. */
template < class T >
void write (CCfits::Column & c, const blitz::Array<T,2>& array, long first_row)
{
  // first check that repeat count is correct
  if (c.repeat() != array.columns())
    throw CCfits::Column::RangeError ("Repeat count not consistent with array size");

  // ensure data is contiguous in memory and that we have a C-array
  blitz::Array<T,2> a (blitz::contiguousArray); 
  if (array.isStorageContiguous() && (a.ordering(0)==1))
    a.reference(array);
  else {
    a.resize(array.shape());
    a=array;
    assert(a.isStorageContiguous());
  }

  assert(a.ordering(0)==1);
  
  c.write(a.dataFirst(), a.size() , a.rows(),  first_row );
}

/** Writes a two-dimensional blitz array to a cfitsio table vector column. */
template < class T >
void write (fitsfile* const fptr, int colnum, const blitz::Array<T,2>& array,
	    long first_row)
{
  int status=0;
  int typecode;
  long repeat;
  long width;
  fits_get_coltype(fptr, colnum, &typecode, &repeat, &width, &status);

  // first check that repeat count is correct
  if (repeat != array.columns())
    throw CCfits::Column::RangeError ("Repeat count not consistent with array size");

  // ensure data is contiguous in memory and that we have a C-array
  blitz::Array<T,2> a (blitz::contiguousArray); 
  if (array.isStorageContiguous() && (a.ordering(0)==1))
    a.reference(array);
  else {
    a.resize(array.shape());
    a=array;
    assert(a.isStorageContiguous());
  }

  assert(a.ordering(0)==1);
  
  fits_write_col (fptr, TDOUBLE, colnum,
		  first_row,1,a.size(), a.dataFirst(), &status);

  checkerr(status);    
}

/** Writes a vector of blitz TinyVectors to a CCfits table vector
    column, starting at row first_row. first_entry and n_rows can be
    specified to only write part of the vector. Default is to write
    the entire vector. */
template < class T, int N >
void write (CCfits::Column & c,
	    const std::vector<blitz::TinyVector<T, N> >& v, 
	    long first_row, long first_entry=0, long n_rows=-1)
{
  if(n_rows==-1)
    n_rows=v.size();
  assert(first_entry+n_rows<=v.size());

  if (c.repeat()!= N)
    throw CCfits::Column::RangeError ("Repeat count not consistent with array size");
  int status=0;

  if(sizeof(blitz::TinyVector<T,N>)!=N*sizeof(T)) {
    // problem - due to alignment constraints, the TinyVectors are
    // padded. We must write one at a time.
    c.parent()->makeThisCurrent();
    CCfits::FITSUtil::MatchType<T> imageType;
    for(int i=0; i<n_rows; ++i)
      fits_write_col (c.parent()->fitsPointer(), imageType(), c.index(),
		      i+first_row,1, N, 
		      (void*)(v[i+first_entry].data()), 
		      &status);
      
    //c.write( const_cast<T*> (v[i+first_entry].data()), N, 1, i+first_row);
  }
  else 
    // TinyVectors are contiguous, we can write all at once 
    c.write( const_cast<T*> (v[first_entry].data()),
	     N*n_rows, n_rows, first_row);

  checkerr(status);
}


/** Writes a 3-dimensional blitz array to a CCfits image
    extension. Because of the need for a temporary to swap storage,
    this is a partial specialization for three-dimensional images,
    which saves memory by writing in slices. */
template <typename T> 
void write (CCfits::ExtHDU& h, const blitz::Array<T, 3>& array)
{
  using namespace blitz;
  const int N=3;
  assert (h.axes() == N);
  blitz::TinyVector<int, N> axis = array.shape();
  for (int i = 0; i < N; ++i)
    if (h.axis(i) != array.extent(i ))
      throw 0;
 
  h.makeThisCurrent();

  // make a temporary 2D array for swapping storage, ensuring it is
  // contiguous in memory
  Array<T, 2> temp (shape (axis [0], axis [1]), 
		    ColumnMajorArray<2> (contiguousData));
  assert(temp.isStorageContiguous());
  // note these may NOT be const because then operator[] returns by value
  TinyVector<long, N> start = 1; 
  TinyVector<long, N> end = axis;
  int status = 0;
  CCfits::FITSUtil::MatchType<T> imageType;
  for (int z = 1; z <= axis [2]; ++z) {
    start [2] = z;
    end [2] = z;
    temp = array (Range::all (), Range::all (), z- 1);
    fits_write_subset (h.fitsPointer(),imageType(), &start[0], &end[0],
		       temp.dataFirst (), &status);
  }
  if (status) {
    fits_report_error (stderr, status);
    throw 1;
  }
 }


/** Writes a 2-dimensional blitz array to a CCfits image
    extension. Because of the need for a temporary to swap storage,
    this is a partial specialization for two-dimensional images, which
    saves memory by writing in slices. */
template <typename T> 
void write (CCfits::ExtHDU& h, const blitz::Array<T, 2>& array)
{
  using namespace blitz;
  const int N=2;
  assert (h.axes() == N);
  blitz::TinyVector<int, N> axis = array.shape();
  for (int i = 0; i < N; ++i)
    assert (h.axis(i) == array.extent(i ));
 
  h.makeThisCurrent();

  // make a temporary 2D array for swapping storage, ensuring it is
  // contiguous in memory
  Array<T, 1> temp (shape (axis [0]), ColumnMajorArray<1> (contiguousData));
  assert(temp.isStorageContiguous());

  TinyVector<long, N> start = 1; 
  TinyVector<long, N> end = axis;
  int status = 0;
  CCfits::FITSUtil::MatchType<T> imageType;
  for (int y = 1; y <= axis [1]; ++y) {
    start [1] = y;
    end [1] = y;
    assert(product(TinyVector<long, N>(end-start+1))==temp.size());
    temp = array (Range::all (), y- 1);
    fits_write_subset (h.fitsPointer(),imageType(), &start[0], &end[0],
		       temp.dataFirst (), &status);
  }

  checkerr(status);    
 }


/** Writes an N-dimensional blitz array to a CCfits image
    extension. FITS images are column-major, so we need to flip the
    storage. */
template <typename HDU, typename T, int N>
void write (HDU& h, const blitz::Array<T,N>& array)
{
  using namespace blitz;

  // check storage order.
  assert (h.axes() == N);
  bool columnmajor=true;
  for (int i = 0; i < N; ++i) {
    assert (h.axis(i) == array.extent(i ));
    columnmajor = columnmajor && (array.ordering(i)==i);
  }

  TinyVector<long, N> start = 1; 
  int status = 0;
  CCfits::FITSUtil::MatchType<T> imageType;
  h.makeThisCurrent();

  // if the array already is column major and contiguous, we don't
  // need the temporary
  if (columnmajor && array.isStorageContiguous()) {
    // write the array data directly using cfitsio
    fits_write_pix(h.fitsPointer(), imageType(), &start[0],
		   array.numElements(), array.dataFirst(), &status);
		   
  }
  else {
    // copy data to a temporary array with proper storage order
    Array<T,N> a (array.shape(), ColumnMajorArray<N> (contiguousData));
    assert(a.isStorageContiguous());
    a = array;

    fits_write_pix(h.fitsPointer(), imageType(), &start[0],
		   a.numElements(), a.dataFirst(), &status);
  }

  checkerr(status);    
}

/** Reads a CCfits image into an N-dimensional blitz array. FITS
    images are column-major, so we need to flip the storage. */
template <typename HDU, typename T, int N>
void read (const HDU& h, blitz::Array<T,N>& array)
{
  using namespace blitz;
  assert (h.axes() == N);

  // check storage order.
  bool columnmajor=true;
  TinyVector<int, N> axis;
  for (int i = 0; i < N; ++i) {
    axis [i] = h.axis(i );
    columnmajor = columnmajor && (array.ordering(i)==i);
  }

  TinyVector<long, N> start = 1; 
  int status = 0;
  CCfits::FITSUtil::MatchType<T> imageType;
  h.makeThisCurrent();

  try {
    // if the array already is column major and contiguous, we don't
    // need the temporary
    if (columnmajor && array.isStorageContiguous()) {
      array.resize(axis);  
      fits_read_pix(h.fitsPointer(), imageType(), &start[0],
		    product(axis), NULL, array.dataFirst(), NULL, &status);
    }
    else {
      // we need to make a temporary
      Array<T,N> a (axis, ColumnMajorArray<N> (contiguousData));
      fits_read_pix(h.fitsPointer(), imageType(), &start[0],
		    product(axis), NULL, a.dataFirst(), NULL, &status);
      array.resize(axis);  
      array=a;
    }
  }
  catch (...) {
    std::cerr << "Array<T,"<< N << "> image read can't allocate output array!"  << std::endl;
    throw;
  }

  checkerr(status);
}

/** Partial specialization (or whatever it's called), reading a
    three-dimensional CCfits data cube into a three-dimensional blitz
    array.  To save memory, we do it in slices instead of allocating a
    completely temporary array. */
template <typename HDU, typename T> 
void read (const CCfits::HDU& h, blitz::Array<T, 3>& array)
{
  using namespace blitz;
  const int N = 3;
  
  assert (h.axes() == N);
  TinyVector<int, N> axis;
  for (int i = 0; i < N; ++i) {
    axis [i] = h.axis(i );
  }
  
  try {
    array.resize(axis );
  }
  catch (...) {
    std::cerr << "Array<T,3> image read can't allocate output array!"  << std::endl;
    throw;
  }

  h.makeThisCurrent();
  Array<T, 2> temp (shape (axis [0], axis [1]), 
		    ColumnMajorArray<2> (contiguousData));
  assert(temp.isStorageContiguous());
  TinyVector<long, N> start = 1; 
  TinyVector<long, N> stride = 1; 
  TinyVector<long, N> end = axis;
  int status = 0, any = 0;
  CCfits::FITSUtil::MatchType<T> imageType;
  for (int z = 1; z <= axis [2]; ++z) {
    start [2] = z;
    end [2] = z;
    fits_read_subset (h.fitsPointer(),imageType(), &start [0], &end [0],
		      &stride [0],0, temp.dataFirst (), &any, &status);
    array (Range::all (), Range::all (), z- 1) = temp;
  }

  checkerr(status);    
}

#endif

